# 滑动多项式

import numpy as np
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import cv2

# 导入图片
binary_warped = mpimg.imread('warped-example.jpg')

def find_lane_pixels(binary_warped):
    # 获取图象下半部分的直方图
    histogram = np.sum(binary_warped[binary_warped.shape[0]//2:,:], axis=0)
    
    # 创建用来输出的图象
    out_img = np.dstack((binary_warped, binary_warped, binary_warped))
    
    # 找到左右两边线的峰值，用在当作左右线的起点
    midpoint = np.int(histogram.shape[0]//2)
    leftx_base = np.argmax(histogram[:midpoint])
    rightx_base = np.argmax(histogram[midpoint:]) + midpoint

    # 定义一些参数
    # 滑动窗口的数量为9
    nwindows = 9

    # 设置滑动窗口的左右边距均为 100
    margin = 100

    # 找到最小像素设置为最近的窗口
    minpix = 50

    # 计算窗口的高度
    window_height = np.int(binary_warped.shape[0]//nwindows)

    # 识别图像中所有非零像素的x,y位置
    nonzero = binary_warped.nonzero()
    nonzeroy = np.array(nonzero[0])
    nonzerox = np.array(nonzero[1])

    # 为窗口更新位置
    leftx_current = leftx_base
    rightx_current = rightx_base

    # 创建列表接收左右车道像素的索引
    left_lane_inds = []
    right_lane_inds = []

    # 一个一个窗口的滑动
    for window in range(nwindows):
        # 计算识别x,y窗口的边接        
        win_y_low = binary_warped.shape[0] - (window + 1) * window_height
        win_y_high = binary_warped.shape[0] - window * window_height
        win_xleft_low = leftx_current - margin
        win_xleft_high = leftx_current + margin
        win_xright_low = rightx_current - margin
        win_xright_high = rightx_current + margin
        
        # 绘制窗口
        # cv2.rectangle openCV中提供的绘制矩形的方法
        # 参数列表 传入图片 (x1, y1) (x2, y2) 这两个坐标分别是矩形的左上角和右下角
        # RGB参数
        # thickness参数 如果为负值 表示填充 正数表示边框厚度
        cv2.rectangle(out_img, (win_xleft_low, win_y_low),
        (win_xleft_high, win_y_high),(0, 255, 0), 2) 
        cv2.rectangle(out_img, (win_xright_low, win_y_low),
        (win_xright_high, win_y_high),(0, 255, 0), 2) 
        
        # 识别窗口内x, y的非零像素
        good_left_inds = ((nonzeroy >= win_y_low) & (nonzeroy < win_y_high) & 
        (nonzerox >= win_xleft_low) &  (nonzerox < win_xleft_high)).nonzero()[0]

        good_right_inds = ((nonzeroy >= win_y_low) & (nonzeroy < win_y_high) & 
        (nonzerox >= win_xright_low) &  (nonzerox < win_xright_high)).nonzero()[0]
        
        # 把找到的非零像素的索引添加进列表里
        left_lane_inds.append(good_left_inds)
        right_lane_inds.append(good_right_inds)
        
        # 如果找到的这个非领像素的长度大于 最小的的像素，呢么下一个窗口的位置会在平均位置
        if len(good_left_inds) > minpix:
            leftx_current = np.int(np.mean(nonzerox[good_left_inds]))
        if len(good_right_inds) > minpix:        
            rightx_current = np.int(np.mean(nonzerox[good_right_inds]))

    # 链接像素的列表
    try:
        left_lane_inds = np.concatenate(left_lane_inds)
        right_lane_inds = np.concatenate(right_lane_inds)
    except ValueError:
        pass

    # 找到左右线的像素位置
    leftx = nonzerox[left_lane_inds]
    lefty = nonzeroy[left_lane_inds] 
    rightx = nonzerox[right_lane_inds]
    righty = nonzeroy[right_lane_inds]

    return leftx, lefty, rightx, righty, out_img


def fit_polynomial(binary_warped):
    # 找到车道像素位置和图片
    leftx, lefty, rightx, righty, out_img = find_lane_pixels(binary_warped)

    # 使用np.polyfit方法拟合多项式
    left_fit = np.polyfit(lefty, leftx, 2)
    right_fit = np.polyfit(righty, rightx, 2)

    # 生成用于绘图的x y
    ploty = np.linspace(0, binary_warped.shape[0]-1, binary_warped.shape[0] )
    try:
        left_fitx = left_fit[0] * ploty ** 2 + left_fit[1] * ploty + left_fit[2]
        right_fitx = right_fit[0] * ploty ** 2 + right_fit[1] * ploty + right_fit[2]
    except TypeError:
        print('没有拟合任何一条线')
        left_fitx = 1*ploty**2 + 1*ploty
        right_fitx = 1*ploty**2 + 1*ploty

    ## 可视化方法 ##
    # 为左右车道设置不同的颜色
    out_img[lefty, leftx] = [255, 0, 0]
    out_img[righty, rightx] = [0, 0, 255]

    # 在左右车到线上绘制函数图象
    plt.plot(left_fitx, ploty, color='yellow')
    plt.plot(right_fitx, ploty, color='yellow')

    return out_img


out_img = fit_polynomial(binary_warped)

plt.imshow(out_img)
plt.show()